#include "randomlib.h"

#include <string.h>
#include <math.h>
#include <stdint.h>

#define MT19937_N 624
#define MT19937_M 397
#define MT19937_MATRIX_A 0x9908b0dfUL
#define MT19937_UPPER_MASK 0x80000000UL
#define MT19937_LOWER_MASK 0x7fffffffUL
#ifndef M_PI
    #define M_PI 3.1415926535897932384626433832795028841971693993751058209749445923078164062862089986280348253421170679
#endif

static void mt19937Twist(mt19937_state* state) {
    uint32_t x;
    uint32_t *mt = state->mt;
    for (int i = 0; i < MT19937_N - MT19937_M; i++) {
        x = (mt[i] & MT19937_UPPER_MASK) | (mt[i+1] & MT19937_LOWER_MASK);
        mt[i] = mt[i + MT19937_M] ^ (x >> 1) ^ (x & 1 ? MT19937_MATRIX_A : 0);
    }
    for (int i = MT19937_N - MT19937_M; i < MT19937_N - 1; i++) {
        x = (mt[i] & MT19937_UPPER_MASK) | (mt[i+1] & MT19937_LOWER_MASK);
        mt[i] = mt[i + MT19937_M - MT19937_N] ^ (x >> 1) ^ (x & 1 ? MT19937_MATRIX_A : 0);
    }
    x = (mt[MT19937_N-1] & MT19937_UPPER_MASK) | (mt[0] & MT19937_LOWER_MASK);
    mt[MT19937_N-1] = mt[MT19937_M-1] ^ (x >> 1) ^ (x & 1 ? MT19937_MATRIX_A : 0);
    state->index = 0;
}

void mt19937Seed(mt19937_state* state, uint32_t seed) {
    state->mt[0] = seed;
    state->index = MT19937_N;

    for (int i = 1; i < MT19937_N; i++) {
        state->mt[i] = 1812433253UL * (state->mt[i-1] ^ (state->mt[i-1] >> 30)) + i;
    }

    mt19937Twist(state);
}

uint32_t mt19937NextInt(mt19937_state* state) {
    if (state->index >= MT19937_N) {
        mt19937Twist(state);
    }

    uint32_t y = state->mt[state->index++];

    y ^= (y >> 11);
    y ^= (y << 7) & 0x9d2c5680UL;
    y ^= (y << 15) & 0xefc60000UL;
    y ^= (y >> 18);

    return y;
}

double mt19937Next(mt19937_state* state) {
    return mt19937NextInt(state) / 4294967296.0;
}

#define MT19937_64_N 312
#define MT19937_64_M 156
#define MT19937_64_MATRIX_A 0xB5026F5AA96619E9ULL
#define MT19937_64_UPPER_MASK 0xFFFFFFFF80000000ULL
#define MT19937_64_LOWER_MASK 0x7FFFFFFFULL

void mt19937_64Seed(mt19937_64_state* state, uint64_t seed) {
    state->mt[0] = seed;
    state->index = MT19937_64_N;

    for (int i = 1; i < MT19937_64_N; i++) {
        state->mt[i] = (6364136223846793005ULL *
                       (state->mt[i-1] ^ (state->mt[i-1] >> 62)) + i);
    }
}

static void mt19937_64Twist(mt19937_64_state* state) {
    uint64_t x;
    uint64_t *mt = state->mt;

    for (int i = 0; i < MT19937_64_N - MT19937_64_M; i++) {
        x = (mt[i] & MT19937_64_UPPER_MASK) | (mt[i+1] & MT19937_64_LOWER_MASK);
        mt[i] = mt[i + MT19937_64_M] ^ (x >> 1) ^ (x & 1 ? MT19937_64_MATRIX_A : 0);
    }

    for (int i = MT19937_64_N - MT19937_64_M; i < MT19937_64_N - 1; i++) {
        x = (mt[i] & MT19937_64_UPPER_MASK) | (mt[i+1] & MT19937_64_LOWER_MASK);
        mt[i] = mt[i + MT19937_64_M - MT19937_64_N] ^ (x >> 1) ^ (x & 1 ? MT19937_64_MATRIX_A : 0);
    }

    x = (mt[MT19937_64_N-1] & MT19937_64_UPPER_MASK) | (mt[0] & MT19937_64_LOWER_MASK);
    mt[MT19937_64_N-1] = mt[MT19937_64_M-1] ^ (x >> 1) ^ (x & 1 ? MT19937_64_MATRIX_A : 0);

    state->index = 0;
}

uint64_t mt19937_64NextInt(mt19937_64_state* state) {
    if (state->index >= MT19937_64_N) {
        mt19937_64Twist(state);
    }

    uint64_t y = state->mt[state->index++];

    // 64位MT19937的tempering变换
    y ^= (y >> 29) & 0x5555555555555555ULL;
    y ^= (y << 17) & 0x71D67FFFEDA60000ULL;
    y ^= (y << 37) & 0xFFF7EEE000000000ULL;
    y ^= (y >> 43);

    return y;
}

double mt19937_64Next(mt19937_64_state* state) {
    return mt19937_64NextInt(state) / 18446744073709551616.0;  // 2^64
}

void xorshift32Seed(xorshift32_state* state, uint32_t seed) {
    state->a = seed ? seed : 1;
}

uint32_t xorshift32NextInt(xorshift32_state* state) {
    state->a ^= state->a << 13;
    state->a ^= state->a >> 17;
    state->a ^= state->a << 5;
    return state->a;
}

double xorshift32Next(xorshift32_state* state) {
    return xorshift32NextInt(state) / 4294967296.0;
}

void lcgSeed(lcg32_state* state, uint32_t seed) {
    state->state = seed;
}

uint32_t lcgNextInt(lcg32_state* state) {
    state->state = state->state * 1103515245 + 12345;
    return (state->state >> 16) & 0x7FFF;
}

double lcgNext(lcg32_state* state) {
    return lcgNextInt(state) / 32768.0;
}

uint32_t pcg32NextInt(pcg32_state* state) {
    uint64_t oldstate = state->state;
    state->state = oldstate * 6364136223846793005ULL + state->inc;

    uint32_t xorshifted = ((oldstate >> 18u) ^ oldstate) >> 27u;
    uint32_t rot = oldstate >> 59u;

    return (xorshifted >> rot) | (xorshifted << ((-rot) & 31));
}

void pcg32Seed(pcg32_state* state, uint32_t seed) {
    state->state = 0;
    state->inc = (seed << 1u) | 1u;
    pcg32NextInt(state);
    state->state += seed;
    pcg32NextInt(state);
}

double pcg32Next(pcg32_state* state) {
    return pcg32NextInt(state) / 4294967296.0;
}

// 生成[minn, maxn)范围内的均匀分布整数
int64_t randint(double n, int64_t minn, int64_t maxn) {
    return minn + (int64_t)(n * (maxn - minn));
}

// 生成[minn, maxn)范围内的均匀分布浮点数
double randfloat(double n, double minn, double maxn) {
    return minn + n * (maxn - minn);
}

// 根据概率生成随机布尔值
bool randbool(double n, double probability) {
    if (probability <= 0.0) return false;
    if (probability >= 1.0) return true;
    return n < probability;
}

// 生成正态分布（高斯分布）随机数
double randnormal(double n1, double n2, double mean, double stddev) {
    // 使用Box-Muller变换
    double z0 = sqrt(-2.0 * log(n1)) * cos(2.0 * M_PI * n2);
    return mean + z0 * stddev;
}

// 生成指数分布随机数
double randexponential(double n, double _lambda) {
    // 使用逆变换法
    return -log(1.0 - n) / _lambda;
}

// 随机打乱数组元素顺序
void shuffle(void* eles, size_t size_ele, size_t n_ele, double (*rng)(void* state), void* state) {
    char* array = (char*)eles;
    char temp_byte;

    // Fisher-Yates洗牌算法
    for (size_t i = n_ele - 1; i > 0; i--) {
        // 生成[0, i]范围内的随机索引
        size_t j = (size_t)(rng(state) * (i + 1));

        if (i != j) {
            // 原地交换元素i和j，逐字节交换
            for (size_t k = 0; k < size_ele; k++) {
                temp_byte = array[i * size_ele + k];
                array[i * size_ele + k] = array[j * size_ele + k];
                array[j * size_ele + k] = temp_byte;
            }
        }
    }
}

void randbytes(uint8_t* output, size_t n_bytes,
               double (*rng)(void* state), void* state) {
    // 使用 uint32_t 来生成 4 个字节
    size_t uint32_count = (n_bytes + 3) / 4; // 向上取整

    for (size_t i = 0; i < uint32_count; i++) {
        uint32_t random_uint32 = (uint32_t)randint(rng(state), 0, UINT32_MAX + 1.0);

        // 计算当前处理的字节位置
        size_t byte_pos = i * 4;

        // 将 uint32_t 分解为 4 个字节
        for (size_t j = 0; j < 4 && byte_pos + j < n_bytes; j++) {
            output[byte_pos + j] = (uint8_t)(random_uint32 >> (j * 8));
        }
    }
}
